# 新增：单张图片测试函数
import math

import torch
import torchvision
from PIL import Image, ImageOps
import matplotlib

from cnn_rotate.utils import img_util

matplotlib.use('TkAgg')
from matplotlib import pyplot as plt

from cnn_rotate import resnet_train, data_load
from torchvision.transforms import functional as F


# 添加字体配置（建议放在脚本开头）
plt.rcParams['font.sans-serif'] = ['SimHei']  # 指定黑体（SimHei）作为中文字体
plt.rcParams['axes.unicode_minus'] = False    # 解决负号显示问题

def test_single_image_rotate(model, image_path, transform, device, num_classes=120):
    model.eval()  # 设置模型为评估模式

    # 读取图片并预处理
    image = Image.open(image_path).convert('RGB')
    # image = data_load.circle_and_crop(image)

    original_image = image.copy()

    image_tensor = img_util.img_rotate(image, 0)

    with torch.no_grad():
        image_transform = transform(image_tensor)
        image_transform = image_transform.unsqueeze(0).to(device)
        outputs = model(image_transform)
        _, predicted_class = torch.max(outputs, 1)

    # 类别转角度（3度/类）
    predicted_angle_deg = predicted_class.item() * (360//num_classes)

    rotated_image_pil = original_image.rotate(-predicted_angle_deg)
    rotated_image_pil = data_load.circle_and_crop(rotated_image_pil)

    plt.figure(figsize=(8, 6))

    # 显示原始图像
    plt.subplot(1, 2, 1)
    plt.imshow(original_image)
    plt.title(f'验证码:', fontsize=12)
    plt.axis('off')
    # 显示旋转后的图像
    plt.subplot(1, 2, 2)
    # plt.imshow(F.to_pil_image(rotated_image_tensor))
    plt.imshow(rotated_image_pil)

    plt.title(f'纠正度数: {predicted_angle_deg}°', fontsize=12)
    plt.axis('off')
    plt.tight_layout()  # 调整子图间距
    plt.show()

    return predicted_angle_deg



def test_single_image(model, image_path, transform, device, num_classes=120, rotate_degree=90, show=False):
    model.eval()  # 设置模型为评估模式

    # 读取图片并预处理
    original_image = Image.open(image_path).convert('RGB')
    image_tensor = img_util.img_rotate(original_image, rotate_degree)

    torchvision.utils.save_image(image_tensor, f"{rotate_degree}.jpg", format='PNG')
    # pil_image = F.to_pil_image(image_tensor)
    # pil_image.save(f"{rotate_degree}.jpg")

    with torch.no_grad():
        image_transform = transform(image_tensor)

        image_transform = image_transform.unsqueeze(0).to(device)  # 添加批次维度并移动到设备

        outputs = model(image_transform)
        _, predicted_class = torch.max(outputs, 1)

    # 类别转角度（3度/类）
    predicted_angle_deg = predicted_class.item() * (360//num_classes)

    # 可视化结果
    if show:
        # 可视化结果（并排显示原始图和旋转图）
        plt.figure(figsize=(8, 6))  # 扩大画布

        # 显示原始图像
        plt.subplot(1, 2, 2)
        plt.imshow(original_image)
        plt.title(f'纠正度数: {predicted_angle_deg}°', fontsize=12)
        plt.axis('off')
        # 显示旋转后的图像
        plt.subplot(1, 2, 1)
        plt.imshow(F.to_pil_image(image_tensor))
        # plt.imshow(image_tensor)

        plt.title('验证码', fontsize=12)
        plt.axis('off')
        plt.tight_layout()  # 调整子图间距
        plt.show()


    return predicted_angle_deg


def test_pic_no_rotate_360():
    test_image_path = '../../dataset/baidu/input/rotated_1748335119.jpg'
    _, val_transform = resnet_train.create_data_transforms()

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = resnet_train.create_model(train=False).to(device)

    model.load_state_dict(torch.load('best_rotation_model.pth', weights_only=True, map_location=device))

    rotate_degree = 146
    predicted_angle = test_single_image(model, test_image_path, val_transform, device, num_classes=120, rotate_degree=rotate_degree, show=True)
    print(f'真实旋转角度: {rotate_degree}---预测旋转角度：{predicted_angle}°')

    # loss_deg = 0
    # for times in range(360):
    #     rotate_degree = times
    #     predicted_angle = test_single_image(model, test_image_path, val_transform, device, num_classes=120, rotate_degree=rotate_degree, show=False)
    #     loss_deg = loss_deg + abs(rotate_degree - predicted_angle)
    #     print(f'真实旋转角度: {rotate_degree}---预测旋转角度：{predicted_angle}°')
    # print(f'平均相差度数: {loss_deg // 360}°')


def test_pic_rotate_1():
    test_image_path = '../../dataset/baidu/test/1749197652.jpg'
    _, val_transform = resnet_train.create_data_transforms()

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = resnet_train.create_model(train=False).to(device)

    model.load_state_dict(torch.load('best_rotation_model.pth', weights_only=True, map_location=device))

    predicted_angle = test_single_image_rotate(model, test_image_path, val_transform, device, num_classes=120)

    print(f'图片: {test_image_path}, 预测旋转角度：{predicted_angle}°')


if __name__ == '__main__':
    # test_pic_no_rotate_360()
    test_pic_rotate_1()
